from .benchmark import Benchmark
import os
import json


class BarBenchmark(Benchmark):
    def __init__(self):
        super().__init__()
        self.pred = dict()
        self.gt = dict()

    def get_area(self, box1, box2):
        if not self.isIntersect(box1, box2):
            return 0, 0, 0

        intersect_area_x1 = box1[0] if box1[0] > box2[0] else box2[0]
        intersect_area_y1 = box1[1] if box1[1] > box2[1] else box2[1]

        intersect_area_x2 = box2[0] + box2[2] if box1[0] + box1[2] > box2[0] + box2[2] else box1[0] + box1[2]
        intersect_area_y2 = box2[1] + box2[3] if box1[1] + box1[3] > box2[1] + box2[3] else box1[1] + box1[3]

        intersect_area = abs(intersect_area_x2 - intersect_area_x1) * abs(intersect_area_y2 - intersect_area_y1)

        box1_area = box1[2] * box1[3]
        box2_area = box2[2] * box2[3]

        return intersect_area, box1_area, box2_area

    @staticmethod
    def isIntersect(box1, box2) -> bool:
        """
        判断两个矩形是否相交
        原理：两个矩形中心点x距离小于两个矩形半宽的和并且两个矩形中心点y距离小于两个矩形半高的和即两个矩形相交
        box1 [x0, y0, w, h]
        box2 [x0, y0, w, h]
        """
        return (abs(2 * box1[0] - 2 * box2[0] + box1[2] - box2[2]) - box1[2] - box2[2]) < 0 \
            and (abs(2 * box1[1] - 2 * box2[1] + box1[3] - box2[3]) - box1[3] - box2[3]) < 0

    def load(self, pred_path: str, gt_path: str):
        """
        加载数据,需要转换成内部格式
        """
        # 读取gt文件
        with open(gt_path, "r") as fp:
            gt = json.load(fp)

        # 赛选gt中饼条形图数据元素
        for key in gt.keys():
            labels = gt[key]
            if "Bar" in labels.keys():
                v = labels["Bar"]
                new_v = []
                for item_index in range(len(v)):
                    item = v[item_index] # 转换成 x,y,w,h 格式
                    new_v.append([item[0], item[1], item[2] - item[0], item[3] - item[1]])
                self.gt[key] = new_v

        # 打开预测结果文件
        with open(pred_path, "r") as fp:
            pred = json.load(fp)

        cnt = 0
        for filename in pred.keys():
            name = filename.split("\\")[-1]
            if "data" not in pred[filename]:
                cnt += 1
                continue

            data = pred[filename]["data"]
            boxes = []
            for i in data:
                if i is None:
                    continue
                for j in i:
                    if j is None:
                        continue
                    # l,t,r,b 转换为 l,t,w,h
                    bbox = j[0]
                    conf = j[1]
                    bbox = bbox[0:2] + [bbox[2] - bbox[0], bbox[3] - bbox[1]]
                    boxes.append([bbox, conf])
            self.pred[name] = boxes
        print(cnt)

    def _eval(self, IOU_threshold: float = 0.5):
        # 类别判断是否有误,IOU小于80%认为是错误的
        bar_recognize = [0, 0]
        # 每一个类别的左上角和宽高的precision和recall
        bar_iou = []
        bar_precision = []
        bar_recall = []
        bar_confidence = []
        time_usage = [-1]

        rec_gt_area = []
        not_rec_gt_area = []

        for gt_key in self.gt.keys():
            if gt_key not in self.pred:
                print(gt_key)
                continue

            gt_value = self.gt[gt_key]
            pred_value = self.pred[gt_key]

            # 最大匹配问题
            for gt_bbox in gt_value:
                max_iou = -1
                max_iou_area = None
                for pred_item in pred_value:
                    pred_bbox = pred_item[0]
                    # 计算IOU
                    intersect_area, gt_area, pred_area = self.get_area(gt_bbox, pred_bbox)
                    if intersect_area == 0:
                        IOU = 0
                    else:
                        IOU = intersect_area / (gt_area + pred_area - intersect_area)
                    if IOU > max_iou:
                        max_iou_area = [intersect_area, gt_area, pred_area, pred_item[1]]
                        max_iou = IOU
                print("MAX_IOU", max_iou)
                # print(IOU)
                if max_iou < IOU_threshold:
                    bar_recognize[0] += 1
                    not_rec_gt_area.append(gt_bbox[-1] * gt_bbox[-2])
                else:
                    rec_gt_area.append(max_iou_area[1])
                    bar_recognize[0] += 1
                    bar_recognize[1] += 1
                    precision = max_iou_area[0] / max_iou_area[1]
                    bar_precision.append(precision)
                    recall = max_iou_area[0] / max_iou_area[2]
                    bar_recall.append(recall)
                    bar_confidence.append(max_iou_area[-1])
                    bar_iou.append(max_iou)

        print(
            f"rec avg area: {sum(rec_gt_area) / len(rec_gt_area)}, not rec avg area: {sum(not_rec_gt_area) / len(not_rec_gt_area)}")
        return bar_recognize, bar_iou, bar_precision, bar_recall, bar_confidence, time_usage

    def eval(self, IOU_threshold: float = 0.5, store_file: str = "cls_benchmark.json"):
        bar_recognize, bar_iou, bar_precision, bar_recall, bar_confidence, time_usage = self._eval(IOU_threshold)
        bar_recognize_rate = bar_recognize[1] / bar_recognize[0]
        bar_iou_average = sum(bar_iou) / len(bar_iou)
        bar_precision_average = sum(bar_precision) / len(bar_precision)
        bar_recall_average = sum(bar_recall) / len(bar_recall)
        bar_confidence_average = sum(bar_confidence) / len(bar_confidence)
        time_usage_per_image_ms = sum(time_usage) / len(time_usage)

        category_name = ["Legend", "ValueAxisTitle", "ChartTitle", "CategoryAxisTitle", "BarPlotArea", "LinePlotArea",
                         "PiePlotArea"]

        print("summary\trecognize_rate\tiou_average\tprecision_average\trecall_average\tconfidence_average")
        print("bar", end="\t")
        print("%.4f" % bar_recognize_rate, end="\t")
        print("%.4f" % bar_iou_average, end="\t")
        print("%.4f" % bar_precision_average, end="\t")
        print("%.4f" % bar_recall_average, end="\t")
        print("%.4f" % bar_confidence_average, end="\t")
        print("")
        print("time usage per image: %.2fms" % time_usage_per_image_ms)

        with open(f"{store_file}", "w") as fp:
            json.dump({
                "raw": {
                    "bar_recognize": bar_recognize,
                    "bar_iou": bar_iou,
                    "bar_precision": bar_precision,
                    "bar_recall": bar_recall,
                    "bar_confidence": bar_confidence,
                    "time_usage": time_usage
                },
                "average": {"recognize_rate": bar_recognize_rate,
                            "iou_average": bar_iou_average,
                            "precision_average": bar_precision_average,
                            "recall_average": bar_recall_average,
                            "confidence_average": bar_confidence_average}
            }, fp)
